import numpy as np

import math
import numpy as np

# ---------- 基本旋转矩阵 ----------
def RotX(a):
    ca, sa = math.cos(a), math.sin(a)
    return np.array([[1, 0, 0],
                     [0, ca, -sa],
                     [0, sa,  ca]], dtype=float)

def RotY(a):
    ca, sa = math.cos(a), math.sin(a)
    return np.array([[ ca, 0, sa],
                     [  0, 1,  0],
                     [-sa, 0, ca]], dtype=float)

def RotZ(a):
    ca, sa = math.cos(a), math.sin(a)
    return np.array([[ca, -sa, 0],
                     [sa,  ca, 0],
                     [ 0,   0, 1]], dtype=float)

# ---------- so(3) <-> SO(3) ----------
def rotvec_to_matrix(r):
    """Rodrigues 公式：旋转向量 -> 旋转矩阵"""
    theta = np.linalg.norm(r)
    if theta < 1e-12:
        return np.eye(3)
    k = r / theta
    K = np.array([[   0, -k[2],  k[1]],
                  [ k[2],    0, -k[0]],
                  [-k[1],  k[0],   0 ]], dtype=float)
    return np.eye(3) + math.sin(theta) * K + (1 - math.cos(theta)) * (K @ K)

def matrix_to_rotvec(R):
    """旋转矩阵 -> 旋转向量（与 UR 的 Rxyz 一致）"""
    # 数值稳健处理
    trace = np.clip((np.trace(R) - 1) / 2.0, -1.0, 1.0)
    theta = math.acos(trace)

    if theta < 1e-12:
        return np.zeros(3)

    # 180 度附近要特别小心
    if abs(math.pi - theta) < 1e-6:
        # 从对角线提取轴方向（参考经典数值技巧）
        R_plus = (R + np.eye(3)) / 2.0
        axis = np.array([math.sqrt(max(R_plus[0,0], 0.0)),
                         math.sqrt(max(R_plus[1,1], 0.0)),
                         math.sqrt(max(R_plus[2,2], 0.0))])
        # 依据非对角项符号修正轴方向
        if R[0,1] < 0: axis[1] = -axis[1]
        if R[0,2] < 0: axis[2] = -axis[2]
        if np.linalg.norm(axis) < 1e-12:
            axis = np.array([1.0, 0.0, 0.0])  # 退化兜底
        axis = axis / np.linalg.norm(axis)
        return axis * theta

    # 通常情况
    axis = np.array([R[2,1] - R[1,2],
                     R[0,2] - R[2,0],
                     R[1,0] - R[0,1]]) / (2.0 * math.sin(theta))
    return axis * theta

# ---------- 本题两函数（约定 R0 = RotX(pi) 为“朝下”的基准） ----------
R0 = RotX(math.pi)  # 工具头面向地面、与世界系对齐的基准姿态

def theta_to_Rxyz(theta):
    """
    输入: theta（弧度），表示在“朝下”基准姿态上绕 工具自身 z 轴 的旋转量
    输出: (rx, ry, rz) —— 可直接给 UR 的 TCP 姿态 Rxyz
    """
    R = R0 @ RotZ(theta)
    r = matrix_to_rotvec(R)
    return float(r[0]), float(r[1]), float(r[2])

def Rxyz_to_theta(rx, ry, rz):
    """
    输入: UR 的 (rx, ry, rz)
    输出: theta（弧度），即在“朝下”基准姿态上绕 工具自身 z 轴 的旋转量
    """
    R = rotvec_to_matrix(np.array([rx, ry, rz], dtype=float))
    # 去掉基准：得到相对旋转（应当接近纯 Rz(theta)）
    R_rel = R0.T @ R
    # 从相对旋转直接读 yaw（注意这是工具系 z 的旋转）
    theta = math.atan2(R_rel[1, 0], R_rel[0, 0])
    return float(theta)


if __name__ == '__main__':
    print(
        # Rxyz_to_z_angle(np.array([0, 3.142, 0]))
        theta_to_Rxyz(0),
        # Rxyz_to_theta(2.596, -1.769, 0)
    )